import matplotlib.pyplot as plt
from matplotlib.pylab import mpl
import numpy as np
from numpy import pi, cos, sin

# 中文显示问题
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

x = [5.5, 2.5, -2., -2., 1., 2.5, 2.5, 2.5, 1., -2., -2., 2.5, 5.5, 2.5, -2., -2., 1., 2.5, 2.5, 2.5, 1., -2., -2., 2.5, 5.5, 2.5, -2., -2., 1., 2.5, 2.5, 2.5, 1., -2., -2., 2.5, 5.5, 2.5, -2., -2., 1., 2.5, 2.5, 2.5, 1., -2., -2., 2.5, 5.5, 2.5, -2., -2., 1., 2.5, 2.5, 2.5, 1., -2., -2., 2.5, 5.5, 2.5, -2., -2., 1., 2.5, 2.5, 2.5, 1., -2., -2., 2.5, 5.5, 2.5, -2., -2., 1., 2.5, 2.5, 2.5, 1., -2., -2., 2.5, 5.5, 2.5, -2., -2., 1., 2.5, 2.5, 2.5, 1., -2., -2., 2.5, 5.5, 2.5, -2., -2., 1., 2.5, 2.5, 2.5, 1., -2., -2., 2.5, 5.5, 2.5, -2., -2., 1., 2.5, 2.5, 2.5, 1., -2., -2., 2.5, 5.5, 2.5, -2., -2., 1., 2.5, 2.5, 2.5, 1., -2., -2., 2.5, 5.5, 2.5, -2., -2., 1., 2.5, 2.5, 2.5, 1., -2., -2., 2.5, 5.5, 2.5, -2., -2., 1., 2.5,
     2.5, 2.5, 1., -2., -2., 2.5, 5.5, 2.5, -2., -2., 1., 2.5, 2.5, 2.5, 1., -2., -2., 2.5, 5.5, 2.5, -2., -2., 1., 2.5, 2.5, 2.5, 1., -2., -2., 2.5, 5.5, 2.5, -2., -2., 1., 2.5, 2.5, 2.5, 1., -2., -2., 2.5, 5.5, 2.5, -2., -2., 1., 2.5, 2.5, 2.5, 1., -2., -2., 2.5, 5.5, 2.5, -2., -2., 1., 2.5, 2.5, 2.5, 1., -2., -2., 2.5, 5.5, 2.5, -2., -2., 1., 2.5, 2.5, 2.5, 1., -2., -2., 2.5, 5.5, 2.5, -2., -2., 1., 2.5, 2.5, 2.5, 1., -2., -2., 2.5, 5.5, 2.5, -2., -2., 1., 2.5, 2.5, 2.5, 1., -2., -2., 2.5, 5.5, 2.5, -2., -2., 1., 2.5, 2.5, 2.5, 1., -2., -2., 2.5, 5.5, 2.5, -2., -2., 1., 2.5, 2.5, 2.5, 1., -2., -2., 2.5, 5.5, 2.5, -2., -2., 1., 2.5, 2.5, 2.5, 1., -2., -2., 2.5, 5.5, 2.5, -2., -2., 1., 2.5, 2.5, 2.5, 1., -2., -2., 2.5]

N = len(x)
freq = np.arange(N)  # 频率counter
fs = N
t = np.arange(0, 1, 1.0 / fs)  # 1秒采样fs个点

X = np.fft.fft(x)  # 离散傅里叶变换

'''
根据STFT公式原理，实现的STFT计算，做了/N的标准化
'''
X2 = np.zeros(N, dtype=np.complex)  # X[n]
for k in range(0, N):  # 0,1,2,...,N-1
    for n in range(0, N):  # 0,1,2,...,N-1
        # X[k] = X[k] + x[n] * np.exp(-2j * pi * k * n / N)
        X2[k] = X2[k] + (1 / N) * x[n] * np.exp(-2j * pi * k * n / N)

# 散点图
plt.subplot(3, 1, 1)
plt.plot(t, x, 'go')
plt.title('散点图')
plt.xlabel('Time(s)')
plt.ylabel('Amplitude')

# 曲线图
plt.subplot(3, 1, 2)
plt.plot(t, x)
plt.title('曲线图')
# plt.xlabel('Time(s)')
plt.ylabel('Amplitude')

freq_half = freq[range(int(N / 2))]  # 前一半频率
X_half = X2[range(int(N / 2))]

# 曲线图
plt.subplot(3, 1, 3)
plt.plot(freq_half, X_half)
plt.title('曲线图')
# plt.xlabel('Time(s)')
plt.ylabel('Amplitude')

fig, ax = plt.subplots(2, 1, figsize=(12, 12))

# 绘制原始时域图像
ax[0].plot(freq, x, label='原始时域信号')
ax[0].set_xlabel('Time (s)')
ax[0].set_ylabel('Amplitude')


ax[1].plot(freq_half, abs(X_half), 'b', label='前N/2个频率')
ax[1].set_xlabel('Freq (Hz)')
ax[1].set_ylabel('Amplitude')
ax[1].set_yticks(np.arange(0, 3))
ax[1].legend()

plt.show()
